# ------------------------------------------------------------------------------
# Copyright (c) Microsoft
# Licensed under the MIT License.
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
# ------------------------------------------------------------------------------

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import time
import logging
import numpy as np
from pathlib import Path
from collections import namedtuple

import torch
import torch.optim as optim
import torch.nn as nn


def concat_batch(batchItem):
    res_list = []
    for item in batchItem:
        tmp = torch.stack(item, dim=0)
        res_list.append(tmp)
    res = torch.cat(res_list)
    return res


def connection_data(a, b):
    assert type(a) == type(b)
    if isinstance(a, np.ndarray):
        return np.concatenate([a, b], axis=0)
    elif isinstance(a, list):
        return a + b
    else:
        ValueError(f'Type {type(a)} data is not supported!')


def merge_dicts(dict_list):
    result = dict_list[0]
    for i in range(1, len(dict_list)):
        dict_item = dict_list[i]
        for key, val in dict_item.items():
            result[key] = connection_data(result[key], val)
    return result



def get_valid_output(outputs, length):
    N = max(length)

    # batch_size = outputs.size(0) // N
    # num_joints = outputs.size(1)

    # outputs = outputs.reshape(batch_size, N, num_joints, outputs.size(-2), outputs.size(-1))
    valid_list = []
    for item, valid in zip(outputs, length):
        valid_output = item[:valid,...]
        valid_list.append(valid_output)

    output = torch.cat(valid_list, dim=0)   #[sum(length), 17, 64, 48]
    return output


def get_valid_output_pure(outputs, length):
    N = max(length)

    batch_size = outputs.size(0) // N
    num_joints = outputs.size(1)

    outputs = outputs.reshape(batch_size, N, num_joints, outputs.size(-2), outputs.size(-1))
    valid_list = []
    for item, valid in zip(outputs, length):
        valid_output = item[:valid,...]
        valid_list.append(valid_output)

    output = torch.cat(valid_list, dim=0)   #[sum(length), 17, 64, 48]
    return output


def create_logger(cfg, cfg_name, global_rank, phase='train'):
    if global_rank == 0:
        root_output_dir = Path(cfg.OUTPUT_DIR)
        # set up logger
        if not root_output_dir.exists():
            print('=> creating {}'.format(root_output_dir))
            root_output_dir.mkdir()

        dataset = cfg.DATASET.DATASET + '_' + cfg.DATASET.HYBRID_JOINTS_TYPE \
            if cfg.DATASET.HYBRID_JOINTS_TYPE else cfg.DATASET.DATASET
        dataset = dataset.replace(':', '_')
        model = cfg.MODEL.NAME
        cfg_name = os.path.basename(cfg_name).split('.')[0]

        final_output_dir = root_output_dir / dataset / model / cfg_name

        print('=> creating {}'.format(final_output_dir))
        final_output_dir.mkdir(parents=True, exist_ok=True)

        time_str = time.strftime('%Y-%m-%d-%H-%M')
        log_file = 'GT_{}_{}_{}_{}.log'.format(cfg.TEST.USE_GT_BBOX, phase, cfg_name, time_str)
        final_log_file = final_output_dir / log_file
        head = '%(asctime)-15s %(message)s'
        logging.basicConfig(filename=str(final_log_file),
                            format=head)
        logger = logging.getLogger()
        logger.setLevel(logging.INFO)
        console = logging.StreamHandler()
        logging.getLogger('').addHandler(console)

        tensorboard_log_dir = Path(cfg.LOG_DIR) / dataset / model / \
            (cfg_name + '_' + time_str)

        print('=> creating {}'.format(tensorboard_log_dir))
        tensorboard_log_dir.mkdir(parents=True, exist_ok=True)
        return logger, str(final_output_dir), str(tensorboard_log_dir)
    else:
        root_output_dir = Path(cfg.OUTPUT_DIR)
        dataset = cfg.DATASET.DATASET + '_' + cfg.DATASET.HYBRID_JOINTS_TYPE \
            if cfg.DATASET.HYBRID_JOINTS_TYPE else cfg.DATASET.DATASET
        dataset = dataset.replace(':', '_')
        model = cfg.MODEL.NAME
        final_output_dir = root_output_dir / dataset / model / cfg_name
        return None, str(final_output_dir), None


def get_optimizer(cfg, parameters):
    optimizer = None
    if cfg.TRAIN.OPTIMIZER == 'sgd':
        optimizer = optim.SGD(
            parameters,
            lr=cfg.TRAIN.LR,
            momentum=cfg.TRAIN.MOMENTUM,
            weight_decay=cfg.TRAIN.WD,
            nesterov=cfg.TRAIN.NESTEROV
        )
    elif cfg.TRAIN.OPTIMIZER == 'adam':
        optimizer = optim.Adam(
            parameters,
            lr=cfg.TRAIN.LR
        )
    elif cfg.TRAIN.OPTIMIZER == 'adamw':
        optimizer = optim.AdamW(
            parameters,
            lr=cfg.TRAIN.LR,
            betas=(0.9,0.999),
            weight_decay=0.01,
        )

    return optimizer


def save_checkpoint(states, is_best, output_dir,
                    filename='checkpoint.pth'):
    torch.save(states, os.path.join(output_dir, filename))
    if is_best and 'state_dict' in states:
        torch.save(states['best_state_dict'],
                   os.path.join(output_dir, 'model_best.pth'))


def get_model_summary(model, *input_tensors, item_length=26, verbose=False):
    """
    :param model:
    :param input_tensors:
    :param item_length:
    :return:
    """

    summary = []

    ModuleDetails = namedtuple(
        "Layer", ["name", "input_size", "output_size", "num_parameters", "multiply_adds"])
    hooks = []
    layer_instances = {}

    def add_hooks(module):

        def hook(module, input, output):
            class_name = str(module.__class__.__name__)

            instance_index = 1
            if class_name not in layer_instances:
                layer_instances[class_name] = instance_index
            else:
                instance_index = layer_instances[class_name] + 1
                layer_instances[class_name] = instance_index

            layer_name = class_name + "_" + str(instance_index)

            params = 0

            if class_name.find("Conv") != -1 or class_name.find("BatchNorm") != -1 or \
               class_name.find("Linear") != -1:
                for param_ in module.parameters():
                    params += param_.view(-1).size(0)

            flops = "Not Available"
            if class_name.find("Conv") != -1 and hasattr(module, "weight"):
                flops = (
                    torch.prod(
                        torch.LongTensor(list(module.weight.data.size()))) *
                    torch.prod(
                        torch.LongTensor(list(output.size())[2:]))).item()
            elif isinstance(module, nn.Linear):
                flops = (torch.prod(torch.LongTensor(list(output.size()))) \
                         * input[0].size(1)).item()

            if isinstance(input[0], list):
                input = input[0]
            if isinstance(output, list):
                output = output[0]

            summary.append(
                ModuleDetails(
                    name=layer_name,
                    input_size=list(input[0].size()),
                    output_size=[0,0],#list(output.size()),
                    num_parameters=params,
                    multiply_adds=flops)
            )

        if not isinstance(module, nn.ModuleList) \
           and not isinstance(module, nn.Sequential) \
           and module != model:
            hooks.append(module.register_forward_hook(hook))

    model.eval()
    model.apply(add_hooks)

    space_len = item_length

    model(*input_tensors)
    for hook in hooks:
        hook.remove()

    details = ''
    if verbose:
        details = "Model Summary" + \
            os.linesep + \
            "Name{}Input Size{}Output Size{}Parameters{}Multiply Adds (Flops){}".format(
                ' ' * (space_len - len("Name")),
                ' ' * (space_len - len("Input Size")),
                ' ' * (space_len - len("Output Size")),
                ' ' * (space_len - len("Parameters")),
                ' ' * (space_len - len("Multiply Adds (Flops)"))) \
                + os.linesep + '-' * space_len * 5 + os.linesep

    params_sum = 0
    flops_sum = 0
    for layer in summary:
        params_sum += layer.num_parameters
        if layer.multiply_adds != "Not Available":
            flops_sum += layer.multiply_adds
        if verbose:
            details += "{}{}{}{}{}{}{}{}{}{}".format(
                layer.name,
                ' ' * (space_len - len(layer.name)),
                layer.input_size,
                ' ' * (space_len - len(str(layer.input_size))),
                layer.output_size,
                ' ' * (space_len - len(str(layer.output_size))),
                layer.num_parameters,
                ' ' * (space_len - len(str(layer.num_parameters))),
                layer.multiply_adds,
                ' ' * (space_len - len(str(layer.multiply_adds)))) \
                + os.linesep + '-' * space_len * 5 + os.linesep

    details += os.linesep \
        + "Total Parameters: {:,}".format(params_sum) \
        + os.linesep + '-' * space_len * 5 + os.linesep
    details += "Total Multiply Adds (For Convolution and Linear Layers only): {:,} GFLOPs".format(flops_sum/(1024**3)) \
        + os.linesep + '-' * space_len * 5 + os.linesep
    details += "Number of Layers" + os.linesep
    for layer in layer_instances:
        details += "{} : {} layers   ".format(layer, layer_instances[layer])

    return details

